#!/usr/bin/env python3
# cython: language_level=3

import random
import numpy as np
import time
import math
import matplotlib.pyplot as plt
from scan_icp import get_ob, icp, match_score, get_trans_matrix, get_rot_matrix, final_score

import rospy
import copy
import sys
sys.path.append("/home/ubuntu/wingbow_ws/src/wingbow_nav/script")
sys.path.append("/home/ubuntu/wingbow_ws/src/wing_scan/script")

import LaserFilter
from LaserFilter import get_obstacle_array
from ScanInfo import ScanInfo
from tf_py2 import euler_from_quaternion
from sensor_msgs.msg import LaserScan
from nav_msgs.msg import Odometry
from threading import Thread

class Plot:
    @staticmethod
    def reset():
        plt.cla()
        plt.axis("equal")
        plt.grid(True)

    @staticmethod
    def plot_data(m, color="ob"):
        plt.plot(m[0, :], m[1, :], color, markersize=2)

'''
简易的搜索匹配点云方法
'''

# ==========真实位置=========== #

demo_running = False
def run_demo(scan):
    global demo_running
    demo_running = True
    print("!! start !!")
    # real_world = get_ob(20, x_len=0.5, width=2.93, length=0.5)
    real_world = np.load("laser_cache.npy")

    z = np.zeros(len(scan[0]))
    z = z+1

    laser_scan = np.vstack([scan, z])
    origin_scan = laser_scan

    Plot.reset()
    Plot.plot_data(real_world, 'or')
    Plot.plot_data(origin_scan)
    plt.pause(0.001)

    best_matrix = np.array([
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]
    ])

    score_cache_0 = 0.0
    score_cache_1 = 10000

    start_time = time.time()
    print(" Start Loop ")
    while (
            abs(score_cache_0-score_cache_1)/score_cache_1>0.3 
            or 
            abs(score_cache_1)>0.03
        ):

        if time.time() - start_time > 5:
            print('[Time Out]')
            break

        score_cache_0 = score_cache_1
        score_cache_1 = match_score(origin_scan, real_world)

        dx, dy, dth = icp(origin_scan, real_world)
        T = get_trans_matrix(dx, dy)
        R = get_rot_matrix(dth)
        best_matrix = R.dot(T.dot(best_matrix))

        # 计算平移距离和旋转角度
        xx = (best_matrix[0,2]/best_matrix[1,0]+best_matrix[1,2]/best_matrix[0,0])/(best_matrix[0,0]/best_matrix[1,0]+best_matrix[1,0]/best_matrix[0,0])
        yy = best_matrix[1,2]/best_matrix[0,0]-xx*best_matrix[1,0]/best_matrix[0,0]
        tth = math.degrees(math.acos(best_matrix[0,0]))
        print("Transform: [%.3f, %.3f, %.2f]"%(xx, yy, tth))

        best_match = R.dot(T.dot(origin_scan))
        Plot.reset()
        Plot.plot_data(real_world, 'or')
        Plot.plot_data(best_matrix.dot(laser_scan), 'og')
        plt.pause(0.001)
        origin_scan = best_match
        
    best_points = final_score(best_match, real_world, 0.05)
    print("There are %d good points in scan."%best_points)
    demo_running = False
    return best_points


class Core:
    def __init__(self):
        self.alive = True
        self.back_sub = None
        self.odom_sub = None
        self.back_info = ScanInfo('back')
        self.back_ob = None
        self.ob_update = False
        self.current_yaw = 0.0

    def init_subs(self):
        self.back_sub = rospy.Subscriber('scan_back', LaserScan, self.back_callback, queue_size=1)
        self.odom_sub = rospy.Subscriber('odom', Odometry, self.odom_callback, queue_size=1)

    def init_scan_info(self, scan, scan_info, trans_r, trans_th, min_ang=-1.658062789, max_ang=1.658062789, threshold=100, max_range=1.0):
        scan_info.scan['ranges'] = list(scan.ranges)
        scan_info.scan['angle_min'] = scan.angle_min
        scan_info.scan['angle_increment'] = scan.angle_increment
        scan_info.scan['intensities'] = list(scan.intensities)
        LaserFilter.init_points(scan_info.scan, threshold=threshold, max_range=max_range, min_angle=min_ang, max_angle=max_ang)
        scan_info.scan = LaserFilter.transform(scan_info.scan, trans_r, trans_th)
        scan_info.DataReady = True

    def odom_callback(self, odom):
        self.current_yaw = euler_from_quaternion(
            [
                odom.pose.pose.orientation.x, 
                odom.pose.pose.orientation.y,
                odom.pose.pose.orientation.z,
                odom.pose.pose.orientation.w,
            ])[2]

    def back_callback(self, scan):
        def reduce_points(scan, max_amount):
            size = len(scan.intensities)
            intesities = list(scan.intensities)
            if size > max_amount:
                intesities.sort()
                return intesities[size-max_amount]
            else:
                return None

        thre = reduce_points(scan, max_amount=200)
        if not thre is None:
            if thre < 80: thre = 80
            pass
        else:
            thre = 80
            
        self.init_scan_info(scan, self.back_info, 0.908, math.pi, min_ang=-math.pi/3, max_ang=math.pi/3, threshold=thre, max_range=8.0)
        if not demo_running:
            self.back_ob = get_obstacle_array(self.back_info.scan).T
            self.ob_update = True

    def icp_task(self):
        rospy.loginfo('ICP DEMO LOOP START.')
        while self.alive:
            time.sleep(0.001)
            if self.ob_update:
                imu_angle = math.degrees(self.current_yaw)
                run_demo(self.back_ob)
                print("IMU Angle is %.2f"%imu_angle)
                self.ob_update = False

    def run(self):
        try:
            rospy.init_node('laser_icp', anonymous=False)
            rospy.loginfo('laser_icp node start.')
            self.init_subs()
            Thread(target=self.icp_task,name='icp_demo').run()
            rospy.spin()
        except BaseException as e:
            rospy.logwarn(e)
        finally:
            self.alive = False

node = Core()

if __name__ == "__main__":
    node.run()

